"""
标准数据模型

本模块定义了市场数据服务的标准数据模型，使用Pydantic实现自动验证和序列化。
这些模型是所有适配器的数据交换格式，确保数据的标准化和类型安全。

核心模型：
1. Bar - K线/蜡烛图数据
2. Tick - 逐笔成交数据
3. Orderbook - 订单簿/深度行情数据

设计原则：
- 使用Pydantic进行自动验证
- 使用Decimal保证价格精度
- 统一时间戳格式
- 提供InfluxDB转换方法

作者: 架构组
创建日期: 2025-10-04
版本: v1.0
"""

from pydantic import BaseModel, Field, field_validator, ConfigDict
from datetime import datetime
from typing import Optional, Dict, Any, List, Tuple
from decimal import Decimal
from enum import Enum


# ==================== 枚举类型 ====================

class DataFrequency(str, Enum):
    """
    数据频率枚举
    
    定义支持的时间周期，用于K线数据查询
    """
    # 秒级
    SECOND_1 = "1s"
    SECOND_5 = "5s"
    SECOND_15 = "15s"
    SECOND_30 = "30s"
    
    # 分钟级
    MINUTE_1 = "1m"
    MINUTE_3 = "3m"
    MINUTE_5 = "5m"
    MINUTE_15 = "15m"
    MINUTE_30 = "30m"
    MINUTE_60 = "60m"
    
    # 小时级
    HOUR_1 = "1h"
    HOUR_2 = "2h"
    HOUR_4 = "4h"
    HOUR_6 = "6h"
    HOUR_12 = "12h"
    
    # 日级
    DAILY = "1d"
    
    # 周级
    WEEKLY = "1w"
    
    # 月级
    MONTHLY = "1M"
    
    # 季度
    QUARTERLY = "1Q"
    
    # 年级
    YEARLY = "1Y"
    
    @classmethod
    def is_valid(cls, frequency: str) -> bool:
        """检查频率是否有效"""
        return frequency in cls._value2member_map_
    
    @classmethod
    def from_string(cls, frequency: str) -> "DataFrequency":
        """
        从字符串创建 DataFrequency 枚举
        
        Args:
            frequency: 频率字符串，如 '1d', '1m', '1h' 等
        
        Returns:
            DataFrequency: 对应的枚举值
        
        Raises:
            ValueError: 如果频率字符串无效
        
        Examples:
            >>> DataFrequency.from_string('1d')
            <DataFrequency.DAILY: '1d'>
            >>> DataFrequency.from_string('1m')
            <DataFrequency.MINUTE_1: '1m'>
        """
        if frequency in cls._value2member_map_:
            return cls._value2member_map_[frequency]
        raise ValueError(f"Invalid frequency: {frequency}. Valid values: {list(cls._value2member_map_.keys())}")


class TickDirection(str, Enum):
    """
    Tick方向枚举
    
    表示逐笔成交是买入还是卖出
    """
    BUY = "buy"       # 主动买入
    SELL = "sell"     # 主动卖出
    UNKNOWN = "unknown"  # 方向未知


class MarketType(str, Enum):
    """
    市场类型枚举
    """
    STOCK = "stock"           # 股票
    FUTURES = "futures"       # 期货
    OPTIONS = "options"       # 期权
    FOREX = "forex"           # 外汇
    CRYPTO = "crypto"         # 加密货币
    INDEX = "index"           # 指数
    FUND = "fund"             # 基金
    BOND = "bond"             # 债券


class Exchange(str, Enum):
    """
    交易所枚举
    """
    # 中国大陆
    SSE = "SSE"         # 上海证券交易所
    SZSE = "SZSE"       # 深圳证券交易所
    SHFE = "SHFE"       # 上海期货交易所
    DCE = "DCE"         # 大连商品交易所
    CZCE = "CZCE"       # 郑州商品交易所
    CFFEX = "CFFEX"     # 中国金融期货交易所
    INE = "INE"         # 上海国际能源交易中心
    
    # 香港
    HKEX = "HKEX"       # 香港交易所
    
    # 美国
    NYSE = "NYSE"       # 纽约证券交易所
    NASDAQ = "NASDAQ"   # 纳斯达克
    AMEX = "AMEX"       # 美国证券交易所
    CME = "CME"         # 芝加哥商品交易所
    CBOT = "CBOT"       # 芝加哥期货交易所
    NYMEX = "NYMEX"     # 纽约商品交易所
    
    # 加密货币
    BINANCE = "BINANCE"       # 币安
    COINBASE = "COINBASE"     # Coinbase
    HUOBI = "HUOBI"           # 火币
    OKEX = "OKEX"             # OKEx


# ==================== 基础模型 ====================

class MarketDataModel(BaseModel):
    """
    市场数据基础模型
    
    所有市场数据模型的基类，包含通用字段和方法
    """
    
    model_config = ConfigDict(
        # Pydantic v2配置
        validate_assignment=True,  # 赋值时验证
        json_encoders={
            Decimal: lambda v: float(v),  # Decimal序列化为float
            datetime: lambda v: v.isoformat()  # datetime序列化为ISO格式
        },
        use_enum_values=True,  # 枚举使用值而不是名称
    )
    
    timestamp: datetime = Field(
        ...,
        description="时间戳（UTC时区）"
    )
    
    symbol: str = Field(
        ...,
        description="标的代码（标准格式）",
        min_length=1,
        max_length=50
    )
    
    @field_validator('symbol')
    @classmethod
    def validate_symbol(cls, v: str) -> str:
        """验证标的代码格式"""
        if not v or v.strip() == "":
            raise ValueError("标的代码不能为空")
        return v.strip().upper()
    
    def to_influx_point(self, measurement: str) -> Dict[str, Any]:
        """
        转换为InfluxDB Point格式
        
        Args:
            measurement: measurement名称
        
        Returns:
            Dict: InfluxDB Point字典
        """
        # 提取tags和fields
        tags = {"symbol": self.symbol}
        fields = {}
        
        for field_name, field_value in self.model_dump().items():
            if field_name in ['timestamp', 'symbol']:
                continue
            if field_value is not None:
                # Decimal转float
                if isinstance(field_value, Decimal):
                    fields[field_name] = float(field_value)
                else:
                    fields[field_name] = field_value
        
        return {
            "measurement": measurement,
            "time": self.timestamp,
            "tags": tags,
            "fields": fields
        }


# ==================== K线数据模型 ====================

class Bar(MarketDataModel):
    """
    K线/蜡烛图数据模型
    
    表示一个时间周期内的价格和成交量信息。
    使用Decimal类型保证价格精度，避免浮点数误差。
    
    属性：
        timestamp: K线时间戳（开盘时间）
        symbol: 标的代码
        open: 开盘价
        high: 最高价
        low: 最低价
        close: 收盘价
        volume: 成交量
        amount: 成交额（可选）
        frequency: 数据频率（如'1m', '1h', '1d'）
        trade_count: 成交笔数（可选）
        vwap: 成交量加权平均价（可选）
    
    示例：
        >>> bar = Bar(
        ...     timestamp=datetime(2025, 10, 4, 9, 30),
        ...     symbol='rb2310.SHFE',
        ...     open=Decimal('3950.0'),
        ...     high=Decimal('3960.0'),
        ...     low=Decimal('3945.0'),
        ...     close=Decimal('3955.0'),
        ...     volume=Decimal('125600'),
        ...     frequency='1m'
        ... )
        >>> print(bar.model_dump_json())
    """
    
    open: Decimal = Field(
        ...,
        description="开盘价",
        ge=0
    )
    
    high: Decimal = Field(
        ...,
        description="最高价",
        ge=0
    )
    
    low: Decimal = Field(
        ...,
        description="最低价",
        ge=0
    )
    
    close: Decimal = Field(
        ...,
        description="收盘价",
        ge=0
    )
    
    volume: Decimal = Field(
        default=Decimal(0),
        description="成交量",
        ge=0
    )
    
    amount: Optional[Decimal] = Field(
        default=None,
        description="成交额",
        ge=0
    )
    
    frequency: str = Field(
        default="1d",
        description="数据频率（如1m, 5m, 1h, 1d）"
    )
    
    trade_count: Optional[int] = Field(
        default=None,
        description="成交笔数",
        ge=0
    )
    
    vwap: Optional[Decimal] = Field(
        default=None,
        description="成交量加权平均价",
        ge=0
    )
    
    @field_validator('high', 'low', 'close')
    @classmethod
    def validate_prices(cls, v: Decimal, info) -> Decimal:
        """验证价格的合理性"""
        if v < 0:
            raise ValueError(f"{info.field_name} 不能为负数")
        return v
    
    def model_post_init(self, __context) -> None:
        """模型初始化后验证"""
        # 验证高低价关系
        if self.high < self.low:
            raise ValueError(f"最高价 ({self.high}) 不能低于最低价 ({self.low})")
        
        # 验证开盘价在高低价范围内
        if not (self.low <= self.open <= self.high):
            raise ValueError(f"开盘价 ({self.open}) 必须在最高价 ({self.high}) 和最低价 ({self.low}) 之间")
        
        # 验证收盘价在高低价范围内
        if not (self.low <= self.close <= self.high):
            raise ValueError(f"收盘价 ({self.close}) 必须在最高价 ({self.high}) 和最低价 ({self.low}) 之间")
    
    def to_influx_point(self) -> Dict[str, Any]:
        """转换为InfluxDB Point格式"""
        return super().to_influx_point(measurement="bars")
    
    @property
    def is_bullish(self) -> bool:
        """是否为阳线"""
        return self.close >= self.open
    
    @property
    def is_bearish(self) -> bool:
        """是否为阴线"""
        return self.close < self.open
    
    @property
    def body_size(self) -> Decimal:
        """实体大小"""
        return abs(self.close - self.open)
    
    @property
    def upper_shadow(self) -> Decimal:
        """上影线长度"""
        return self.high - max(self.open, self.close)
    
    @property
    def lower_shadow(self) -> Decimal:
        """下影线长度"""
        return min(self.open, self.close) - self.low
    
    @property
    def range(self) -> Decimal:
        """振幅（高低价差）"""
        return self.high - self.low


# ==================== Tick数据模型 ====================

class Tick(MarketDataModel):
    """
    Tick逐笔成交数据模型
    
    表示单笔成交记录，包含价格、数量、方向等信息。
    
    属性：
        timestamp: 成交时间
        symbol: 标的代码
        price: 成交价格
        volume: 成交数量
        direction: 成交方向（买入/卖出/未知）
        bid_price_1: 买一价（可选）
        ask_price_1: 卖一价（可选）
        bid_volume_1: 买一量（可选）
        ask_volume_1: 卖一量（可选）
        last_price: 最新价（可选）
        total_volume: 总成交量（可选）
        total_amount: 总成交额（可选）
        open_interest: 持仓量（期货专用，可选）
    
    示例：
        >>> tick = Tick(
        ...     timestamp=datetime.now(),
        ...     symbol='BTC-USD',
        ...     price=Decimal('45000.50'),
        ...     volume=Decimal('0.5'),
        ...     direction=TickDirection.BUY
        ... )
    """
    
    price: Decimal = Field(
        ...,
        description="成交价格",
        ge=0
    )
    
    volume: Decimal = Field(
        ...,
        description="成交数量",
        ge=0
    )
    
    direction: TickDirection = Field(
        default=TickDirection.UNKNOWN,
        description="成交方向"
    )
    
    # 买卖盘信息
    bid_price_1: Optional[Decimal] = Field(
        default=None,
        description="买一价",
        ge=0
    )
    
    ask_price_1: Optional[Decimal] = Field(
        default=None,
        description="卖一价",
        ge=0
    )
    
    bid_volume_1: Optional[Decimal] = Field(
        default=None,
        description="买一量",
        ge=0
    )
    
    ask_volume_1: Optional[Decimal] = Field(
        default=None,
        description="卖一量",
        ge=0
    )
    
    # 其他信息
    last_price: Optional[Decimal] = Field(
        default=None,
        description="最新价",
        ge=0
    )
    
    total_volume: Optional[Decimal] = Field(
        default=None,
        description="总成交量",
        ge=0
    )
    
    total_amount: Optional[Decimal] = Field(
        default=None,
        description="总成交额",
        ge=0
    )
    
    open_interest: Optional[Decimal] = Field(
        default=None,
        description="持仓量（期货）",
        ge=0
    )
    
    # 扩展字段，用于存储额外信息
    extra_fields: Optional[Dict[str, Any]] = Field(
        default=None,
        description="扩展字段，用于存储额外的市场数据信息"
    )
    
    def to_influx_point(self) -> Dict[str, Any]:
        """转换为InfluxDB Point格式"""
        return super().to_influx_point(measurement="ticks")
    
    @property
    def spread(self) -> Optional[Decimal]:
        """买卖价差"""
        if self.bid_price_1 is not None and self.ask_price_1 is not None:
            return self.ask_price_1 - self.bid_price_1
        return None
    
    @property
    def mid_price(self) -> Optional[Decimal]:
        """中间价"""
        if self.bid_price_1 is not None and self.ask_price_1 is not None:
            return (self.bid_price_1 + self.ask_price_1) / Decimal(2)
        return None


# ==================== 订单簿数据模型 ====================

class OrderbookLevel(BaseModel):
    """
    订单簿单个档位
    
    表示买盘或卖盘的一个价格档位
    """
    
    model_config = ConfigDict(
        json_encoders={
            Decimal: lambda v: float(v)
        }
    )
    
    price: Decimal = Field(
        ...,
        description="价格",
        ge=0
    )
    
    volume: Decimal = Field(
        ...,
        description="数量",
        ge=0
    )
    
    order_count: Optional[int] = Field(
        default=None,
        description="订单数量",
        ge=0
    )


class Orderbook(MarketDataModel):
    """
    订单簿/深度行情数据模型
    
    表示某一时刻的买卖盘深度信息。
    
    属性：
        timestamp: 时间戳
        symbol: 标的代码
        bids: 买盘列表，按价格降序排列
        asks: 卖盘列表，按价格升序排列
        sequence: 序列号（用于检测丢包，可选）
    
    示例：
        >>> orderbook = Orderbook(
        ...     timestamp=datetime.now(),
        ...     symbol='BTC-USD',
        ...     bids=[
        ...         OrderbookLevel(price=Decimal('45000'), volume=Decimal('1.5')),
        ...         OrderbookLevel(price=Decimal('44999'), volume=Decimal('2.0'))
        ...     ],
        ...     asks=[
        ...         OrderbookLevel(price=Decimal('45001'), volume=Decimal('1.0')),
        ...         OrderbookLevel(price=Decimal('45002'), volume=Decimal('1.8'))
        ...     ]
        ... )
        >>> print(f"最优买价: {orderbook.best_bid_price}")
        >>> print(f"最优卖价: {orderbook.best_ask_price}")
    """
    
    bids: List[OrderbookLevel] = Field(
        default_factory=list,
        description="买盘列表（按价格降序）"
    )
    
    asks: List[OrderbookLevel] = Field(
        default_factory=list,
        description="卖盘列表（按价格升序）"
    )
    
    sequence: Optional[int] = Field(
        default=None,
        description="序列号（用于检测丢包）"
    )
    
    @field_validator('bids')
    @classmethod
    def validate_bids(cls, v: List[OrderbookLevel]) -> List[OrderbookLevel]:
        """验证买盘按价格降序排列"""
        if len(v) > 1:
            for i in range(len(v) - 1):
                if v[i].price < v[i + 1].price:
                    raise ValueError("买盘必须按价格降序排列")
        return v
    
    @field_validator('asks')
    @classmethod
    def validate_asks(cls, v: List[OrderbookLevel]) -> List[OrderbookLevel]:
        """验证卖盘按价格升序排列"""
        if len(v) > 1:
            for i in range(len(v) - 1):
                if v[i].price > v[i + 1].price:
                    raise ValueError("卖盘必须按价格升序排列")
        return v
    
    @property
    def best_bid_price(self) -> Optional[Decimal]:
        """最优买价"""
        return self.bids[0].price if self.bids else None
    
    @property
    def best_ask_price(self) -> Optional[Decimal]:
        """最优卖价"""
        return self.asks[0].price if self.asks else None
    
    @property
    def best_bid_volume(self) -> Optional[Decimal]:
        """最优买量"""
        return self.bids[0].volume if self.bids else None
    
    @property
    def best_ask_volume(self) -> Optional[Decimal]:
        """最优卖量"""
        return self.asks[0].volume if self.asks else None
    
    @property
    def spread(self) -> Optional[Decimal]:
        """买卖价差"""
        if self.best_bid_price is not None and self.best_ask_price is not None:
            return self.best_ask_price - self.best_bid_price
        return None
    
    @property
    def mid_price(self) -> Optional[Decimal]:
        """中间价"""
        if self.best_bid_price is not None and self.best_ask_price is not None:
            return (self.best_bid_price + self.best_ask_price) / Decimal(2)
        return None
    
    @property
    def bid_depth(self) -> int:
        """买盘深度"""
        return len(self.bids)
    
    @property
    def ask_depth(self) -> int:
        """卖盘深度"""
        return len(self.asks)
    
    @property
    def total_bid_volume(self) -> Decimal:
        """买盘总量"""
        return sum(level.volume for level in self.bids)
    
    @property
    def total_ask_volume(self) -> Decimal:
        """卖盘总量"""
        return sum(level.volume for level in self.asks)
    
    def to_influx_point(self) -> Dict[str, Any]:
        """转换为InfluxDB Point格式"""
        # 订单簿通常不直接存储到InfluxDB（数据量大）
        # 而是提取关键指标存储
        fields = {
            "best_bid_price": float(self.best_bid_price) if self.best_bid_price else None,
            "best_ask_price": float(self.best_ask_price) if self.best_ask_price else None,
            "best_bid_volume": float(self.best_bid_volume) if self.best_bid_volume else None,
            "best_ask_volume": float(self.best_ask_volume) if self.best_ask_volume else None,
            "spread": float(self.spread) if self.spread else None,
            "mid_price": float(self.mid_price) if self.mid_price else None,
            "bid_depth": self.bid_depth,
            "ask_depth": self.ask_depth,
            "total_bid_volume": float(self.total_bid_volume),
            "total_ask_volume": float(self.total_ask_volume),
        }
        
        # 移除None值
        fields = {k: v for k, v in fields.items() if v is not None}
        
        return {
            "measurement": "orderbook",
            "time": self.timestamp,
            "tags": {"symbol": self.symbol},
            "fields": fields
        }


# ==================== 辅助函数 ====================

def parse_symbol(symbol: str) -> Tuple[str, Optional[Exchange]]:
    """
    解析标的代码，提取代码和交易所
    
    Args:
        symbol: 标的代码，格式如 'rb2310.SHFE' 或 '000001.SZ'
    
    Returns:
        Tuple[str, Optional[Exchange]]: (代码, 交易所)
    
    示例:
        >>> code, exchange = parse_symbol('rb2310.SHFE')
        >>> print(code)  # 'rb2310'
        >>> print(exchange)  # Exchange.SHFE
    """
    if '.' in symbol:
        code, exchange_str = symbol.split('.', 1)
        try:
            exchange = Exchange(exchange_str.upper())
        except ValueError:
            exchange = None
        return code, exchange
    return symbol, None


def build_symbol(code: str, exchange: Exchange) -> str:
    """
    构建标准格式的标的代码
    
    Args:
        code: 代码
        exchange: 交易所
    
    Returns:
        str: 标准格式的标的代码
    
    示例:
        >>> symbol = build_symbol('rb2310', Exchange.SHFE)
        >>> print(symbol)  # 'rb2310.SHFE'
    """
    return f"{code}.{exchange.value}"

